import time

import numpy as np
import torch
from torch.utils.data import Dataset
from torchvision import transforms
import torch.nn.functional as F
from torch.distributions.normal import Normal
from modules import *

# 加载数据集
signal_path = 'signal.npy'
data_x = EmWaveLoader(signal_path) # float64,cpu
print("数据加载完毕，其中一份数据为：\n",data_x.data[0,:])

# 先测试1000-20-1001的简易模型
device = torch.device('cuda:0')
# vae=VAE(1001,1001,10).to(device)
vae=VAE(1001,1001,10)
print('\nvae 模型结构:\n',vae)

# opt = torch.optim.Adam(vae.parameters(), lr=0.005, amsgrad=True)
BEST_LOSS=999999999
# 模型训练过程
for epoch in range(1, 2):
    print("Epoch {}:".format(epoch))
    train(vae,data_x)
    cur_loss = test(vae,data_x)  # 当前模型的损失函数计算

    # 保存最优模型
    if cur_loss <= BEST_LOSS:
        BEST_LOSS = cur_loss
        LAST_SAVED = epoch
        print("Saving model!")
        torch.save(vae.state_dict(), 'models/{}_vae.pt'.format('test'))
    else:
        print("Not saving model! Last saved: {}".format(LAST_SAVED))

    # generate_reconstructions()
    # generate_samples()


